package com.huawei.dli.flink.sye.source;

import java.util.Properties;

import com.huawei.dli.flink.sye.utils.LoginUtil;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class KafkaKerbosSource<T> extends FlinkKafkaConsumer<T> {
  private static final Logger logger = LoggerFactory.getLogger(KafkaKerbosSource.class);

  private String tmpDir = "/tmp/kafkaSource.";
  private String principalName;

  public KafkaKerbosSource(String topic, DeserializationSchema<T> valueDeserializer,
                           Properties properties, String userPrincipal) {
    super(topic, valueDeserializer, properties);
    this.principalName = userPrincipal;
  }

  @Override
  public void open(Configuration parameters) throws Exception {
    securityPrepare();
    super.open(parameters);
  }

  @Override
  public void close() throws Exception {
    super.close();
    try {
      LoginUtil.deleteFile(tmpDir, "user.keytab");
      LoginUtil.deleteFile(tmpDir, "krb5.conf");
    } catch (Exception e) {
      logger.error("Failed to close source.", e);
    }
  }

  private void securityPrepare() throws Exception {
    String userKeytabFile = LoginUtil.createFile(tmpDir, "user.keytab", this.getClass());
    String krb5File = LoginUtil.createFile(tmpDir, "krb5.conf", this.getClass());
    LoginUtil.setKrb5Config(krb5File);
    LoginUtil.setJaasConf(LoginUtil.Module.KAFKA.getName(), principalName, userKeytabFile);
    sun.security.krb5.Config.refresh();
  }
}
